from functools import wraps
from flask import request, jsonify, g
from models import User, UserRole
import jwt
import os

def token_required(f):
    """需要有效token的装饰器"""
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None
        
        # 从请求头获取token
        if 'Authorization' in request.headers:
            auth_header = request.headers['Authorization']
            try:
                token = auth_header.split(" ")[1]  # Bearer <token>
            except IndexError:
                return jsonify({'error': 'Invalid token format'}), 401
        
        if not token:
            return jsonify({'error': 'Token is missing'}), 401
        
        try:
            # 解码token
            data = jwt.decode(token, os.getenv('SECRET_KEY', 'your-secret-key'), algorithms=['HS256'])
            current_user = User.query.get(data['user_id'])
            if not current_user or not current_user.is_active:
                return jsonify({'error': 'Invalid token or user inactive'}), 401
            g.current_user = current_user
        except jwt.ExpiredSignatureError:
            return jsonify({'error': 'Token has expired'}), 401
        except jwt.InvalidTokenError:
            return jsonify({'error': 'Invalid token'}), 401
        
        return f(*args, **kwargs)
    return decorated

def permission_required(permission):
    """需要特定权限的装饰器"""
    def decorator(f):
        @wraps(f)
        @token_required
        def decorated(*args, **kwargs):
            if not g.current_user.has_permission(permission):
                return jsonify({
                    'error': 'Insufficient permissions',
                    'required_permission': permission
                }), 403
            return f(*args, **kwargs)
        return decorated
    return decorator

def admin_required(f):
    """需要管理员权限的装饰器"""
    @wraps(f)
    @token_required
    def decorated(*args, **kwargs):
        if not g.current_user.is_admin():
            return jsonify({'error': 'Admin access required'}), 403
        return f(*args, **kwargs)
    return decorated

def super_admin_required(f):
    """需要超级管理员权限的装饰器"""
    @wraps(f)
    @token_required
    def decorated(*args, **kwargs):
        if not g.current_user.is_super_admin():
            return jsonify({'error': 'Super admin access required'}), 403
        return f(*args, **kwargs)
    return decorated

def role_required(required_roles):
    """需要特定角色的装饰器"""
    if not isinstance(required_roles, list):
        required_roles = [required_roles]
    
    def decorator(f):
        @wraps(f)
        @token_required
        def decorated(*args, **kwargs):
            if g.current_user.role not in required_roles:
                return jsonify({
                    'error': 'Insufficient role',
                    'required_roles': [role.value for role in required_roles],
                    'current_role': g.current_user.role.value
                }), 403
            return f(*args, **kwargs)
        return decorated
    return decorator

def model_access_required(f):
    """检查模型访问权限的装饰器"""
    @wraps(f)
    @token_required
    def decorated(*args, **kwargs):
        # 从请求中获取模型名称
        model_name = None
        if request.is_json:
            model_name = request.json.get('model')
        elif request.form:
            model_name = request.form.get('model')
        
        if model_name and not g.current_user.can_access_model(model_name):
            return jsonify({
                'error': 'Model access denied',
                'model': model_name,
                'allowed_models': g.current_user.get_permissions().get('models', [])
            }), 403
        
        return f(*args, **kwargs)
    return decorated

def check_request_limit(f):
    """检查请求限制的装饰器"""
    @wraps(f)
    @token_required
    def decorated(*args, **kwargs):
        if not g.current_user.can_make_request():
            return jsonify({
                'error': 'Weekly request limit exceeded',
                'current_count': g.current_user.get_weekly_request_count(),
                'limit': g.current_user.get_effective_weekly_limit()
            }), 429
        
        return f(*args, **kwargs)
    return decorated

def get_user_permissions():
    """获取当前用户权限的辅助函数"""
    if hasattr(g, 'current_user'):
        return g.current_user.get_permissions()
    return None

def check_concurrent_requests(f):
    """检查并发请求限制的装饰器"""
    @wraps(f)
    @token_required
    def decorated(*args, **kwargs):
        permissions = g.current_user.get_permissions()
        max_concurrent = permissions.get('max_concurrent_requests', 1)
        
        if max_concurrent == -1:  # 无限制
            return f(*args, **kwargs)
        
        # 这里可以添加实际的并发请求检查逻辑
        # 例如使用Redis或内存缓存来跟踪活跃请求
        
        return f(*args, **kwargs)
    return decorated